import json
import uuid
import asyncio
from datetime import datetime, timedelta
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth import get_user_model
from django.utils import timezone
from drf_spectacular_websocket.decorators import extend_ws_schema
from .models import AIConversation, AIMessage
from .models_streaming import StreamSession
from .ai_service import AIService
from .ws_serializers import (
    StreamStartRequestSerializer, StreamStartResponseSerializer,
    StreamStopRequestSerializer, StreamStoppedResponseSerializer,
    PingRequestSerializer, PongResponseSerializer,
    StreamChunkResponseSerializer, StreamCompleteResponseSerializer,
    StreamErrorResponseSerializer, ErrorResponseSerializer
)
import logging

logger = logging.getLogger(__name__)


class ChatStreamConsumer(AsyncWebsocketConsumer):
    """AI聊天流式输出WebSocket消费者
    
    WebSocket连接URL: ws://host:port/ws/ai-chat/conversations/{conversation_id}/stream/
    
    支持的消息类型:
    - start_stream: 开始AI流式生成
    - stop_stream: 停止AI流式生成 
    - ping: 心跳检测
    
    服务器推送消息类型:
    - stream_start: 流式生成开始
    - stream_chunk: 流式内容块
    - stream_complete: 流式生成完成
    - stream_error: 流式生成错误
    - stream_stopped: 流式生成停止
    - error: 通用错误
    - pong: 心跳响应
    """
    
    @extend_ws_schema(
        type='receive',
        summary='流式内容块推送',
        description='服务器推送AI生成的内容块到客户端',
        responses=StreamChunkResponseSerializer
    )
    async def send_stream_chunk(self):
        """文档用：流式内容块推送"""
        pass
    
    @extend_ws_schema(
        type='receive',
        summary='流式生成完成通知',
        description='服务器通知客户端流式生成已完成',
        responses=StreamCompleteResponseSerializer
    )
    async def send_stream_complete(self):
        """文档用：流式生成完成通知"""
        pass
    
    @extend_ws_schema(
        type='receive',
        summary='流式生成错误通知',
        description='服务器通知客户端流式生成出现错误',
        responses=StreamErrorResponseSerializer
    )
    async def send_stream_error(self):
        """文档用：流式生成错误通知"""
        pass
    
    @extend_ws_schema(
        type='send',
        summary='心跳检测',
        description='客户端发送心跳检测消息',
        request=PingRequestSerializer,
        responses=PongResponseSerializer
    )
    async def handle_ping(self):
        """文档用：心跳检测处理"""
        pass
    
    async def connect(self):
        """WebSocket连接建立"""
        print(f'=== WebSocket连接建立开始 ===')
        self.conversation_id = self.scope['url_route']['kwargs']['conversation_id']
        self.user = self.scope['user']
        print(f'对话ID: {self.conversation_id}')
        print(f'用户: {self.user}')
        print(f'用户已认证: {self.user.is_authenticated}')
        
        # 验证用户是否已认证
        if not self.user.is_authenticated:
            print('❌ 用户未认证，关闭连接')
            await self.close()
            return
        
        # 验证对话是否存在且用户有权限访问
        try:
            print('🔍 验证对话权限...')
            self.conversation = await self.get_conversation(self.conversation_id, self.user)
            if not self.conversation:
                print('❌ 对话不存在或无权限访问')
                await self.close()
                return
            print(f'✅ 对话验证成功: {self.conversation.title}')
        except Exception as e:
            print(f'❌ 获取对话失败: {e}')
            logger.error(f"获取对话失败: {e}")
            await self.close()
            return
        
        # 加入房间组
        self.room_group_name = f'chat_stream_{self.conversation_id}'
        print(f'🏠 加入房间组: {self.room_group_name}')
        await self.channel_layer.group_add(
            self.room_group_name,
            self.channel_name
        )
        
        await self.accept()
        print('🎉 WebSocket连接已接受')
        print(f'=== WebSocket连接建立完成 ===')
        logger.info(f"用户 {self.user.username} 连接到对话 {self.conversation_id} 的流式输出")
    
    async def disconnect(self, close_code):
        """WebSocket连接断开"""
        print(f'=== WebSocket连接断开 ===')
        print(f'关闭代码: {close_code}')
        
        # 离开房间组
        if hasattr(self, 'room_group_name'):
            print(f'🏠 离开房间组: {self.room_group_name}')
            await self.channel_layer.group_discard(
                self.room_group_name,
                self.channel_name
            )
        
        # 清理可能存在的流式会话
        if hasattr(self, 'current_stream_session'):
            print('🧹 清理流式会话')
            await self.cleanup_stream_session(self.current_stream_session)
        
        print(f'=== WebSocket断开完成 ===')
        logger.info(f"用户断开连接，关闭代码: {close_code}")
    
    async def receive(self, text_data):
        """接收来自WebSocket的消息"""
        print(f'📨 收到WebSocket消息: {text_data[:100]}...' if len(text_data) > 100 else f'📨 收到WebSocket消息: {text_data}')
        try:
            data = json.loads(text_data)
            message_type = data.get('type')
            print(f'📋 消息类型: {message_type}')
            
            if message_type == 'start_stream':
                print('🚀 处理开始流式输出请求')
                await self.handle_start_stream(data)
            elif message_type == 'stop_stream':
                print('🛑 处理停止流式输出请求')
                await self.handle_stop_stream(data)
            elif message_type == 'ping':
                print('🏓 处理心跳检测')
                await self.send_json({'type': 'pong'})
            else:
                print(f'❓ 未知消息类型: {message_type}')
                logger.warning(f"未知消息类型: {message_type}")
                
        except json.JSONDecodeError:
            logger.error("接收到无效的JSON数据")
            await self.send_json({
                'type': 'error',
                'message': '无效的JSON格式'
            })
        except Exception as e:
            logger.error(f"处理消息时发生错误: {e}")
            await self.send_json({
                'type': 'error',
                'message': '处理消息时发生错误'
            })
    
    @extend_ws_schema(
        type='send',
        summary='开始AI流式生成',
        description='发送消息开始AI流式回复生成',
        request=StreamStartRequestSerializer,
        responses={
            200: StreamStartResponseSerializer,
            400: ErrorResponseSerializer,
        }
    )
    async def handle_start_stream(self, data):
        """处理开始流式输出请求"""
        try:
            # 提取必要参数
            message_content = data.get('message', '').strip()
            model = data.get('model', 'gpt-3.5-turbo')
            role = data.get('role', 'user')
            
            if not message_content:
                await self.send_json({
                    'type': 'error',
                    'message': '消息内容不能为空'
                })
                return
            
            # 创建流式会话记录
            stream_id = str(uuid.uuid4())
            self.current_stream_session = await self.create_stream_session(
                stream_id, self.conversation, self.user, message_content, model
            )
            
            # 通知开始流式输出
            await self.send_json({
                'type': 'stream_start',
                'stream_id': stream_id
            })
            
            # 开始AI流式生成
            await self.start_ai_stream(stream_id, message_content, model, role)
            
        except Exception as e:
            logger.error(f"开始流式输出时发生错误: {e}")
            await self.send_json({
                'type': 'error',
                'message': f'开始流式输出失败: {str(e)}'
            })
    
    @extend_ws_schema(
        type='send',
        summary='停止AI流式生成',
        description='停止当前正在进行的AI流式回复生成',
        request=StreamStopRequestSerializer,
        responses={
            200: StreamStoppedResponseSerializer,
        }
    )
    async def handle_stop_stream(self, data):
        """处理停止流式输出请求"""
        stream_id = data.get('stream_id')
        if hasattr(self, 'current_stream_session') and self.current_stream_session:
            await self.stop_stream_session(self.current_stream_session, 'user_stopped')
            await self.send_json({
                'type': 'stream_stopped',
                'stream_id': stream_id
            })
    
    async def start_ai_stream(self, stream_id, message_content, model, role):
        """开始AI流式生成"""
        try:
            # 更新流式会话状态
            await self.update_stream_session_status(
                self.current_stream_session, 'streaming'
            )
            
            # 获取对话历史
            conversation_history = await self.get_conversation_history(self.conversation)
            
            # 预加载AI模型和提供商信息
            conversation_with_relations = await self.get_conversation_with_relations(self.conversation_id)
            if not conversation_with_relations:
                raise Exception("无法获取对话详细信息")
            
            # 调用AI服务进行流式生成
            ai_service = AIService(conversation_with_relations)
            
            async def stream_callback(chunk_data):
                """流式数据回调函数"""
                try:
                    content = chunk_data.get('content', '')
                    is_final = chunk_data.get('is_final', False)
                    tokens = chunk_data.get('tokens', 0)
                    
                    # 更新流式会话的当前内容
                    await self.update_stream_session_content(
                        self.current_stream_session, content, tokens
                    )
                    
                    # 发送流式数据到前端
                    await self.send_json({
                        'type': 'stream_chunk',
                        'stream_id': stream_id,
                        'content': content,
                        'tokens': tokens,
                        'is_final': is_final
                    })
                    
                    if is_final:
                        # 获取AI消息ID（从流式回调数据中获取）
                        ai_message_id = chunk_data.get('ai_message_id')
                        
                        # 完成流式会话
                        if ai_message_id:
                            ai_message = await self.get_message_by_id(ai_message_id)
                            await self.complete_stream_session(
                                self.current_stream_session, content, tokens, ai_message
                            )
                        else:
                            await self.complete_stream_session_without_message(
                                self.current_stream_session, content, tokens
                            )
                        
                        # 通知流式输出完成
                        await self.send_json({
                            'type': 'stream_complete',
                            'stream_id': stream_id,
                            'message_id': ai_message_id,
                            'total_tokens': tokens
                        })
                        
                except Exception as e:
                    logger.error(f"流式回调处理错误: {e}")
                    await self.send_json({
                        'type': 'error',
                        'message': f'流式处理错误: {str(e)}'
                    })
            
            # 执行流式生成 - 现在使用上下文管理器
            await ai_service.stream_chat(
                user_message_content=message_content,
                user=self.user,
                stream_callback=stream_callback,
                include_context=True,
                context_limit=10,
                max_context_tokens=120000,
                parent_message_id=None
            )
            
        except Exception as e:
            logger.error(f"AI流式生成错误: {e}")
            await self.fail_stream_session(self.current_stream_session, str(e))
            await self.send_json({
                'type': 'stream_error',
                'stream_id': stream_id,
                'message': f'AI生成失败: {str(e)}'
            })
    
    async def send_json(self, data):
        """发送JSON数据到WebSocket"""
        await self.send(text_data=json.dumps(data, ensure_ascii=False))
    
    # 数据库操作方法
    @database_sync_to_async
    def get_conversation(self, conversation_id, user):
        """获取对话"""
        try:
            return AIConversation.objects.get(id=conversation_id, owner=user)
        except AIConversation.DoesNotExist:
            return None
    
    @database_sync_to_async
    def get_conversation_with_relations(self, conversation_id):
        """获取对话及相关外键信息"""
        try:
            return AIConversation.objects.select_related(
                'ai_model', 'ai_model__provider', 'ai_role'
            ).get(id=conversation_id)
        except AIConversation.DoesNotExist:
            return None
    
    @database_sync_to_async
    def create_stream_session(self, stream_id, conversation, user, content, model):
        """创建流式会话记录"""
        return StreamSession.objects.create(
            stream_id=stream_id,
            conversation=conversation,
            user=user,
            status='pending',
            started_at=timezone.now(),
            expires_at=timezone.now() + timedelta(hours=1)  # 1小时后过期
        )
    
    @database_sync_to_async
    def create_message(self, conversation, user, content, role):
        """创建消息记录"""
        # 根据role设置ai_role字段
        ai_role = None if role == 'user' else conversation.ai_role
        
        return AIMessage.objects.create(
            conversation=conversation,
            ai_role=ai_role,
            content=content
        )
    
    @database_sync_to_async
    def update_stream_session_status(self, session, status):
        """更新流式会话状态"""
        session.status = status
        if status == 'streaming' and not session.started_at:
            session.started_at = timezone.now()
        session.save(update_fields=['status', 'started_at'])
    
    @database_sync_to_async
    def update_stream_session_content(self, session, content, tokens):
        """更新流式会话内容"""
        session.current_content = content
        session.current_tokens = tokens
        session.save(update_fields=['current_content', 'current_tokens'])
    
    @database_sync_to_async
    def complete_stream_session(self, session, final_content, total_tokens, message):
        """完成流式会话"""
        session.status = 'completed'
        session.current_content = final_content
        session.current_tokens = total_tokens
        session.message = message
        session.completed_at = timezone.now()
        if session.started_at:
            duration = timezone.now() - session.started_at
            session.duration_ms = int(duration.total_seconds() * 1000)
        session.save()
    
    @database_sync_to_async
    def fail_stream_session(self, session, error_message):
        """失败流式会话"""
        session.status = 'failed'
        session.error_message = error_message
        session.completed_at = timezone.now()
        if session.started_at:
            duration = timezone.now() - session.started_at
            session.duration_ms = int(duration.total_seconds() * 1000)
        session.save()
    
    @database_sync_to_async
    def stop_stream_session(self, session, reason):
        """停止流式会话"""
        session.status = 'timeout' if reason == 'timeout' else 'completed'
        session.completed_at = timezone.now()
        if session.started_at:
            duration = timezone.now() - session.started_at
            session.duration_ms = int(duration.total_seconds() * 1000)
        session.save()
    
    @database_sync_to_async
    def cleanup_stream_session(self, session):
        """清理流式会话"""
        if session and session.status in ['pending', 'streaming']:
            session.status = 'timeout'
            session.completed_at = timezone.now()
            session.save()
    
    @database_sync_to_async
    def get_conversation_history(self, conversation):
        """获取对话历史"""
        messages = AIMessage.objects.filter(conversation=conversation).select_related('ai_role').order_by('created_at')
        return [
            {
                'role': 'user' if msg.ai_role_id is None else 'assistant',
                'content': msg.content,
                'created_at': msg.created_at.isoformat()
            }
            for msg in messages
        ]
    
    @database_sync_to_async
    def get_message_by_id(self, message_id):
        """根据ID获取消息"""
        try:
            return AIMessage.objects.get(id=message_id)
        except AIMessage.DoesNotExist:
            return None
    
    @database_sync_to_async
    def complete_stream_session_without_message(self, session, final_content, total_tokens):
        """完成流式会话（无AI消息）"""
        session.status = 'completed'
        session.current_content = final_content
        session.current_tokens = total_tokens
        session.completed_at = timezone.now()
        if session.started_at:
            duration = timezone.now() - session.started_at
            session.duration_ms = int(duration.total_seconds() * 1000)
        session.save()